import pandas as pd
from pathlib import Path
import warnings

# --- 配置参数 ---
# cBioPortal 突变数据文件的路径，与 all_step1_mutation_features_and_ids.py 中使用的路径一致
CBIOPORTAL_INPUT_CSV_PATH = "/gpfs/flash/home/yzq/project/model/Mutabert/data/raw/cbioportal_data/cbioportal_mutations_grch38_converted.csv"

# 输出文件路径，用于保存独特的 studyId 列表
# 建议放在一个合适的输出目录下，例如 VOCAB_OUTPUT_DIR 或 CLINICAL_DATA_OUTPUT_DIR
OUTPUT_COHORTS_FILE_PATH = "/gpfs/flash/home/yzq/project/model/SomaticMutaGPT/src/scripts/cbioportal_unique_study_ids.txt"

def save_cbioportal_cohorts_to_file(csv_path: str, output_file_path: str):
    """
    加载 cBioPortal 突变数据，统计其中独特的 studyId，并将其保存到指定文件中。

    Args:
        csv_path (str): cBioPortal 突变数据 CSV 文件的路径。
        output_file_path (str): 保存独特 studyId 列表的文本文件路径。
    """
    csv_file = Path(csv_path)
    output_file = Path(output_file_path)

    if not csv_file.exists():
        print(f"错误：输入文件未找到：{csv_file}")
        return

    # 确保输出目录存在
    output_file.parent.mkdir(parents=True, exist_ok=True)

    print(f"正在加载 cBioPortal 数据集：{csv_file}")
    try:
        df = pd.read_csv(csv_file, low_memory=False)
        print(f"数据加载成功。原始数据形状：{df.shape}")
    except Exception as e:
        print(f"加载 CSV 文件时发生错误：{e}")
        return

    # 检查 'studyId' 列是否存在
    if 'studyId' not in df.columns:
        print("错误：数据集中未找到 'studyId' 列。无法统计和保存队列。")
        print(f"可用列：{df.columns.tolist()}")
        return

    # 获取唯一的 studyId
    unique_study_ids = df['studyId'].unique()
    
    # 统计数量
    num_unique_study_ids = len(unique_study_ids)

    print(f"\n--- cBioPortal 数据集队列统计与保存 ---")
    print(f"cBioPortal 数据集中包含 {num_unique_study_ids} 个独特的队列 (studyId)。")
    
    # 将唯一的 studyId 保存到文件
    try:
        with open(output_file, 'w', encoding='utf-8') as f:
            for study_id in unique_study_ids:
                f.write(f"{study_id}\n")
        print(f"独特的 studyId 列表已成功保存到：{output_file}")
    except Exception as e:
        print(f"保存 studyId 列表到文件时发生错误：{e}")

if __name__ == '__main__':
    save_cbioportal_cohorts_to_file(CBIOPORTAL_INPUT_CSV_PATH, OUTPUT_COHORTS_FILE_PATH)